import torch
import torch.nn.functional as F
import matplotlib.pyplot as plt
import random

def build_dataset(words):
    X, Y = [], []
    for w in words:
            #print(w)
            context = [0] * block_size
            for ch in w + '.':
                idx = stoi[ch]
                X.append(context)
                Y.append(idx)
                #print(''.join(itos[i] for i in context), '--->', itos[idx])
                context = context[1:] + [idx]

    X =   torch.tensor(X)
    Y =   torch.tensor(Y)
    print(X.shape, Y.shape)
    return X, Y

def draw_historgram(Layers):
    plt.figure(figsize=(20,4))
    legends = []
    for i, layer in enumerate(Layers):
        if isinstance(layer, Tanh):
            t = layer.out
            print("layer %d (%10s): mean %+.2f, std %.2f, saturated: %.2f" % (i, layer.__class__, t.mean(), t.std(), (t > 0.97).float().mean() * 100)) 
            hy, hx = torch.histogram(t, density=True)
            plt.plot(hx[:-1].tolist(), hy.tolist())
            legends.append(f'layer {i}')
    plt.legend(legends)
    plt.title("activation histogram")
    plt.show()

def draw_histogram_grad(Layers):
    plt.figure(figsize=(20,4))
    legends = []
    for i, layer in enumerate(Layers):
        if isinstance(layer, Tanh):
            t = layer.out.grad
            print("layer %d (%10s): mean %+.2f, std %.2f" % (i, layer.__class__, t.mean(), t.std()))
            hy, hx = torch.histogram(t, density=True)
            plt.plot(hx[:-1].tolist(), hy.tolist())
            legends.append(f'layer {i}')
    plt.legend(legends)
    plt.title("gradient histogram")
    plt.show()

def draw_2dim_histogram(parameters):
    plt.figure(figsize=(20,4))
    legends = []
    for i, p in enumerate(parameters):
        if p.ndim == 2:
            t = p.grad
            print("weight(%10s): mean %+f, std %.e, grand:data ratio %e" % (tuple(p.shape), t.mean(), t.std(), t.std() / p.std()))
            hy, hx = torch.histogram(t, density=True)
            plt.plot(hx[:-1].tolist(), hy.tolist())
            legends.append(f'{i} {tuple(p.shape)}')
    plt.legend(legends)
    plt.title("gradient histogram")
    plt.show()

def draw_ud_histogram(ud, parameters):
    plt.figure(figsize=(20,4))
    legends = []
    for i, p in enumerate(parameters):
        plt.plot([ud[j][i] for j in range(len(ud))])
        legends.append('param %d' % i)
    plt.plot([0, len(ud)], [-3, -3], 'k') # these ratios shold be ~1e03
    plt.legend(legends)
    plt.title("gradient histogram")
    plt.show()

class Linear:
    def __init__(self, fan_in, fan_out, bias=True, generator=None):
        self.W = torch.randn(fan_in, fan_out, generator=generator) / fan_in**0.5
        self.b = torch.zeros(fan_out) if bias else None
    def __call__(self, x):
        self.out = x @ self.W
        if self.b is not None:
            self.out += self.b
        return self.out
    def parameters(self):
        return [self.W] if self.b is None else [self.W, self.b]
    
class BatchNorm1d:

    def __init__(self, dim, eps=1e-5, momentum=0.1):
        self.eps = eps
        self.momentum = momentum
        self.training = True
        #  parameters (trained with backprop)
        self.gamma = torch.ones(dim)
        self.beta = torch.zeros(dim)
        #  buffers (trained with a running 'momentum')
        self.running_mean = torch.zeros(dim)
        self.running_var = torch.ones(dim)
    def __call__(self, x):
        if self.training:
            xmean = x.mean(0, keepdim=True)
            xvar = x.var(0, keepdim=True)
        else:
            xmean = self.running_mean
            xvar = self.running_var
        xhat = (x - xmean) / torch.sqrt(xvar + self.eps)
        self.out = self.gamma * xhat + self.beta
        if self.training:
            with torch.no_grad():
                self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * xmean
                self.running_var = (1 - self.momentum) * self.running_var + self.momentum * xvar

        return self.out
    def parameters(self):
        return [self.gamma, self.beta]
    
class Tanh:
    def __init__(self):
        self.out = None
    def __call__(self, x):
        self.out = torch.tanh(x)
        return self.out
    def parameters(self):
        return []
    
def inference(model):
    for i in range(20):
        out = []
        context = [0] * block_size()
        while True:
            logits = model(torch.tensor([context]))
            ix = torch.multinomial(F.softmax(logits, dim=1), num_samples=1).item()
            out.append(ix)
            context = context[1:] + [ix]
            if ix == 0:
                break
        print(''.join(itos[i] for i in out))
if __name__ == '__main__':
    words = open('names.txt', 'r').read().splitlines()
    print(len(words))
    chars = sorted(list(set(''.join(words))))
    chars = ['.'] + chars
    stoi = {s:i for i,s in enumerate(chars)}
    itos = {s:i for i,s in stoi.items()}    
    vocab_size = len(stoi)
    n_embed = 10  # dimension of the embedding
    n_hidden = 200 # number of neurons in hidden layers
    block_size = 3
    n = len(words)
    n1 = int (0.8*n)
    n2 = int (0.9*n)
    X, Y = build_dataset(words[:n1])
    Xval, Yval = build_dataset(words[n1:n2])
    Xtest, Ytest = build_dataset(words[n2:])
    g = torch.Generator().manual_seed(2147483647)
    C=  torch.randn((vocab_size, n_embed), generator=g)
    Layers = [
        Linear(n_embed * block_size, n_hidden, generator=g), BatchNorm1d(n_hidden), Tanh(),
        Linear(          n_hidden, n_hidden, generator=g), BatchNorm1d(n_hidden), Tanh(),
        Linear(          n_hidden, n_hidden, generator=g), BatchNorm1d(n_hidden), Tanh(),
        Linear(          n_hidden, n_hidden, generator=g), BatchNorm1d(n_hidden), Tanh(),
        Linear(          n_hidden, n_hidden, generator=g), BatchNorm1d(n_hidden), Tanh(),
        Linear(          n_hidden, vocab_size, generator=g), BatchNorm1d(vocab_size),
    ]

    with torch.no_grad():
        Layers[-1].gamma *= 0.1 #  last layer make less confident
        for layer in Layers[:-1]:
            if isinstance(layer, Linear):
                layer.W *= 5/3
    
    parameters = [C] + [p for layer in Layers for p in layer.parameters()]
    print("number of parameters in total:",sum(p.nelement() for p in parameters)) #  number of parameters in total
    for p in parameters:
        p.requires_grad = True

    max_steps = 200000  
    batch_size = 32
    lossi = []
    ud = []
    for i in range(max_steps):
        #  minibatch construct
        idx = torch.randint(0, X.shape[0], (batch_size,)) 
        Xb = X[idx]
        Yb = Y[idx]

        #  forward pass
        emb = C[Xb]  #  embedding lookup
        x = emb.view(emb.shape[0], -1)
        print(x.shape)
        for layer in Layers:
            x = layer(x)
        loss = F.cross_entropy(x, Yb) #  loss function

        #   backward pass
        for layer in Layers:
            layer.out.retain_grad() 
        
        for p in parameters:
            p.grad = None

        loss.backward()

        lr = 0.1 if i < 100000 else 0.01
        for p in parameters:
            p.data += -lr * p.grad

        if i % 10000 == 0:
            print(f'iter {i:7d}/{max_steps:7d}: {loss.item():.4f}')
        with torch.no_grad():
            ud.append([(lr*p.grad.std()/p.data.std()).log10().item() for p in parameters])
        lossi.append(loss.log10().item())

        if i > 1000:
            break
    draw_historgram(Layers)
    draw_histogram_grad(Layers)
    draw_2dim_histogram(parameters)
    draw_ud_histogram(ud, parameters)

